#!/usr/bin/env python

import optparse
import os
import sys

import numpy
import tables

from charmicat import charmicat.BadCommandLineError

def main():
    '''
    ptpercent.py [opts] FRAC file.h5 [file2.h5 [...]]

    Given each input file `file.h5', write an output file containing some fraction FRAC
    of the events in the input file.
    '''

    op = optparse.OptionParser(usage=main.__doc__)

    op.add_option('-o', dest='output', default='.',
                  help='Leave output in directory named NAME (will be created if non-existent)',
                  metavar='NAME')
    op.add_option('-r', '--random', action='store_true', default=False, dest='random',
                  help='Randomly sample input instead of taking fraction at beginning of file')
    op.add_option("--nt", dest='nt', default='nt',
                  help='Use table located at fp.root.NT', metavar='NT')
    op.add_option("-z", "--compress", dest='compress', action='store_true', default=False,
                  help='Compress output file')

    
    opts, args = op.parse_args()

    try:
        frac = float(args.pop(0))
    except:
        msg = "Please supply a floating-point number between 0 and 1 as the first argument to this script."
        raise charmicat.BadCommandLineError(msg, -1)

    if len(args) == 0:
        msg = "Please supply at least one input file."
        raise charmicat.BadCommandLineError(msg, -2)

    if os.access(opts.output, os.F_OK):
        if not os.path.isdir(opts.output):
            msg = "Output path `" + opts.output + "' already exists as something other than a directory.\n"
            msg += "Please specify another output path or remove the existing one.\n"
            raise charmicat.BadCommandLineError(msg, -3)
    else:        
        os.mkdir(opts.output)


    for i, ipath in enumerate(args):
        sys.stdout.write("{0}/{1} files complete\r".format(i, len(args)))
        sys.stdout.flush()

        fpin  = tables.openFile(ipath)

        try:
            ntin = eval("fpin.root." + opts.nt)
        except tables.NoSuchNodeError:
            sys.stderr.write('\n')
            sys.stderr.write("No node `fp.root." + opts.nt + "' in file " + ipath + "\n")
            sys.stderr.write("Continuing...\n");
            sys.stderr.flush()
            fpin.close()
            continue

        if round(len(ntin) * frac) == 0:
            sys.stderr.write('\n')
            sys.stderr.write("{0:.1f}% of the rows in `{1}' yields 0 rows ({2} total)\n".format(100*frac,
                                                                                                ipath,
                                                                                                len(ntin)))
                                                                                            
            sys.stderr.write("Continuing...\n");
            sys.stderr.flush()
            fpin.close()
            continue
            
        
        opath =  '{0}.{1}pct.h5'.format(os.path.basename(ipath).rsplit('.', 1)[0], 
                                        int(100*frac))

        opath = os.path.join(opts.output, opath)

        fpout = tables.openFile(opath, 'w')
        
        comment = "Generated by `pyt_percent.py'; {0}% of {1}".format(int(frac*100), os.path.abspath(ipath))
        
        if opts.compress:
            filters = tables.Filters(complevel=1, complib='lzo', shuffle=True)
        else:
            filters = None

        ntout = fpout.createTable(fpout.root, opts.nt, ntin.coldescrs, comment, filters=filters)

        totrows  = len(ntin)
        fracrows = int(round(totrows * frac))
        
        if opts.random:
            indices = numpy.random.randint(0, len(ntin), fracrows)
            ntout.append(ntin.readCoordinates(indices))
        else:
            ntout.append(ntin.read(0, stop=fracrows))
        
        ntout.flush()

        fpout.close()
        fpin.close()


    print "{0}/{1} files complete\r".format(len(args), len(args))

    return 0
        

        
            

    
    

if __name__ == "__main__":
    try:
        sys.exit(main())
    except charmicat.BadCommandLineError, e:
        sys.stderr.write(str(e) + "\n")
        sys.stderr.write("See `pyt_ten_percent.py -h' for more.\n")
        sys.stderr.flush()
        sys.exit(e.retcode)
